# people_downloader_ui_flow_no_mouse.py
import hashlib
import os
import threading
import time
import tkinter as tk
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from io import BytesIO
from tkinter import filedialog, messagebox, ttk
from urllib.parse import urlencode

import requests
from PIL import Image, ImageTk, ImageOps

URL = "https://thispersondoesnotexist.com/"
TARGET_SIZE = 100 * 1024  # 100KB
PREVIEW_SIZE = (200, 200)  # 预览缩略图尺寸 WxH（200x200）

session = requests.Session()
session.headers.update({
    "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
    "Cache-Control": "no-cache",
    "Pragma": "no-cache",
    "User-Agent": "Mozilla/5.0"
})


def fetch_one():
    """下载一张图片并返回 {'raw': bytes, 'img': PIL.Image} 或抛出异常"""
    ts_query = urlencode({"_": int(time.time() * 1000)})
    full_url = f"{URL}?{ts_query}"
    resp = session.get(full_url, timeout=20)
    resp.raise_for_status()
    if "image" not in resp.headers.get("Content-Type", ""):
        raise ValueError("非图片响应")
    raw = resp.content
    img = Image.open(BytesIO(raw)).convert("RGB")
    return {"raw": raw, "img": img}


def compress_image(img: Image.Image):
    """把 PIL.Image 压缩到 <= TARGET_SIZE（尽量），返回 bytes, quality"""
    quality = 95
    step = 5
    out = BytesIO()
    while quality > 5:
        out.seek(0)
        out.truncate()
        img.save(out, format="JPEG", quality=quality, optimize=True)
        if out.tell() <= TARGET_SIZE:
            break
        quality -= step
    return out.getvalue(), quality


class App:
    def __init__(self, root):
        self.root = root
        self.root.title("获取测试人脸数据 By:Qiuweihuang")
        self.images = []  # 列表：{raw, img, thumb_tk, var, md5, frame}
        self.hash_set = set()  # md5 去重
        self.item_frames = []  # 对应 GUI 元素列表

        # 顶部工具栏
        toolbar = tk.Frame(root)
        toolbar.pack(fill="x", padx=6, pady=6)

        tk.Label(toolbar, text="数量:").pack(side="left")
        self.count_var = tk.IntVar(value=8)
        tk.Entry(toolbar, textvariable=self.count_var, width=6).pack(side="left")

        tk.Label(toolbar, text="线程:").pack(side="left", padx=(8, 0))
        self.thread_var = tk.IntVar(value=4)
        tk.Entry(toolbar, textvariable=self.thread_var, width=6).pack(side="left")

        btn_fetch = tk.Button(toolbar, text="获取图片", command=self.on_fetch)
        btn_fetch.pack(side="left", padx=8)
        self.btn_fetch = btn_fetch

        tk.Button(toolbar, text="选择保存目录", command=self.choose_dir).pack(side="left")
        self.save_dir = os.getcwd()
        self.path_label = tk.Label(toolbar, text=f"保存目录: {self.save_dir}")
        self.path_label.pack(side="left", padx=8)

        self.save_orig_var = tk.BooleanVar(value=True)
        self.save_comp_var = tk.BooleanVar(value=True)
        tk.Checkbutton(toolbar, text="保存原图", variable=self.save_orig_var).pack(side="left")
        tk.Checkbutton(toolbar, text="保存压缩图", variable=self.save_comp_var).pack(side="left")

        # 全选 / 反选 按钮
        tk.Button(toolbar, text="全选", command=self.select_all).pack(side="right", padx=4)
        tk.Button(toolbar, text="反选", command=self.invert_selection).pack(side="right", padx=4)
        tk.Button(toolbar, text="保存选中", command=self.save_selected, bg="green", fg="white").pack(side="right",
                                                                                                     padx=6)

        # 进度条（加载动画）
        self.progress = ttk.Progressbar(root, mode="indeterminate")
        self.progress.pack(fill="x", padx=6, pady=(0, 6))
        self.progress.stop()
        self.progress.pack_forget()  # 初始隐藏

        # 预览区：使用 Canvas + 内部 frame（用 grid 动态换行）
        preview_container = tk.Frame(root)
        preview_container.pack(fill="both", expand=True)
        self.canvas = tk.Canvas(preview_container)
        self.scroll_y = tk.Scrollbar(preview_container, orient="vertical", command=self.canvas.yview)
        self.inner_frame = tk.Frame(self.canvas)

        self.canvas.create_window((0, 0), window=self.inner_frame, anchor="nw")
        self.canvas.configure(yscrollcommand=self.scroll_y.set)
        self.canvas.pack(side="left", fill="both", expand=True)
        self.scroll_y.pack(side="right", fill="y")

        self.inner_frame.bind("<Configure>", lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all")))
        self.canvas.bind("<Configure>", lambda e: self.reflow_items())

        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        tk.Label(root, textvariable=self.status_var, anchor="w").pack(fill="x", padx=6, pady=(4, 6))

    def choose_dir(self):
        folder = filedialog.askdirectory()
        if folder:
            self.save_dir = folder
            self.path_label.config(text=f"保存目录: {self.save_dir}")

    def set_busy(self, busy=True):
        """切换加载动画和按钮可用性"""
        if busy:
            self.progress.pack(fill="x", padx=6, pady=(0, 6))
            self.progress.start(10)
            self.btn_fetch.config(state="disabled")
        else:
            self.progress.stop()
            self.progress.pack_forget()
            self.btn_fetch.config(state="normal")

    def on_fetch(self):
        """点击 获取 图片"""
        try:
            count = int(self.count_var.get())
            threads = int(self.thread_var.get())
            if count <= 0: raise ValueError("数量必须大于0")
            if threads <= 0: raise ValueError("线程必须大于0")
        except Exception as e:
            messagebox.showerror("参数错误", str(e))
            return

        # 清空现有预览
        self.images.clear()
        self.hash_set.clear()
        for w in self.inner_frame.winfo_children():
            w.destroy()
        self.item_frames.clear()

        self.set_busy(True)
        self.status_var.set(f"开始获取 {count} 张（去重）...")

        # 后台线程管理下载逻辑（不直接操作 UI）
        thread = threading.Thread(target=self._background_fetch, args=(count, threads), daemon=True)
        thread.start()

    def _background_fetch(self, desired_count, threads):
        """后台下载循环：提交小批次，遇到重复继续尝试，直到收集到 desired_count 或达到最大尝试"""
        max_attempts = max(desired_count * 6, 50)
        attempts = 0

        with ThreadPoolExecutor(max_workers=threads) as ex:
            unique_collected = 0
            while unique_collected < desired_count and attempts < max_attempts:
                batch = min(threads, desired_count - unique_collected)
                futures = [ex.submit(fetch_one) for _ in range(batch)]
                attempts += batch
                for f in as_completed(futures):
                    if unique_collected >= desired_count:
                        break
                    try:
                        data = f.result()
                        md5 = hashlib.md5(data["raw"]).hexdigest()
                        # 去重判断
                        if md5 in self.hash_set:
                            continue
                        self.hash_set.add(md5)
                        unique_collected = len(self.hash_set)
                        # 将 UI 添加任务交给主线程
                        self.root.after(0, self.add_image_to_ui, data, md5)
                        time.sleep(0.03)
                    except Exception as e:
                        self.root.after(0, lambda msg=str(e): self._append_status(f"下载错误: {msg}"))
        self.root.after(0, self._fetch_done, len(self.hash_set), attempts)

    def _append_status(self, txt):
        self.status_var.set(txt)

    def _fetch_done(self, got, attempts):
        self.set_busy(False)
        self.status_var.set(f"完成：获取到 {got} 张（尝试 {attempts} 次，去重后）")

    def add_image_to_ui(self, data, md5):
        """在主线程创建缩略和界面控件（缩略 200x200），单击放大"""
        # 生成固定 200x200 预览（居中裁剪）
        try:
            thumb = ImageOps.fit(data["img"], PREVIEW_SIZE, Image.LANCZOS)
        except Exception:
            thumb = data["img"].copy()
            thumb.thumbnail(PREVIEW_SIZE, Image.LANCZOS)
            bg = Image.new("RGB", PREVIEW_SIZE, (255, 255, 255))
            bg.paste(thumb, ((PREVIEW_SIZE[0] - thumb.width) // 2, (PREVIEW_SIZE[1] - thumb.height) // 2))
            thumb = bg

        tk_thumb = ImageTk.PhotoImage(thumb)

        # 每个图片项的容器
        item = tk.Frame(self.inner_frame, bd=2, relief="ridge")
        lbl = tk.Label(item, image=tk_thumb, cursor="hand2")
        lbl.image = tk_thumb  # 防止回收
        lbl.pack(side="top")
        # 显示 md5 前 8 位作为 id
        id_label = tk.Label(item, text=f"ID: {md5[:8]}")
        id_label.pack(side="top", anchor="w", padx=4, pady=(4, 0))

        var = tk.BooleanVar(value=True)
        cb = tk.Checkbutton(item, text="保存", variable=var)
        cb.pack(side="left", padx=6, pady=6)

        # 单击打开原始大图预览
        def on_click(evt, img=data["img"]):
            self.open_large_preview(img)

        lbl.bind("<Button-1>", on_click)

        # 保存内部记录
        record = {
            "raw": data["raw"],
            "img": data["img"],
            "thumb_tk": tk_thumb,
            "var": var,
            "md5": md5,
            "frame": item
        }
        self.images.append(record)
        self.item_frames.append(item)

        # 布局重排
        self.reflow_items()

    def reflow_items(self):
        """根据 canvas 可见宽度，动态计算每行列数并用 grid 放置项以实现换行"""
        pad_x = 8
        pad_y = 8
        item_w = PREVIEW_SIZE[0] + 16  # 单项宽度估计（含边与间距）
        avail = self.canvas.winfo_width()
        if avail <= 0:
            self.root.after(100, self.reflow_items)
            return
        cols = max(1, avail // item_w)
        for w in self.item_frames:
            w.grid_forget()
        for idx, w in enumerate(self.item_frames):
            r = idx // cols
            c = idx % cols
            w.grid(row=r, column=c, padx=pad_x, pady=pad_y, sticky="n")
        self.inner_frame.update_idletasks()
        self.canvas.configure(scrollregion=self.canvas.bbox("all"))

    def open_large_preview(self, pil_image: Image.Image):
        """弹窗显示原始大图（适配屏幕大小）"""
        top = tk.Toplevel(self.root)
        top.title("大图预览")
        w, h = pil_image.size
        screen_w = top.winfo_screenwidth()
        screen_h = top.winfo_screenheight()
        max_w = int(screen_w * 0.8)
        max_h = int(screen_h * 0.8)
        disp = pil_image.copy()
        if w > max_w or h > max_h:
            disp.thumbnail((max_w, max_h), Image.LANCZOS)
        tk_img = ImageTk.PhotoImage(disp)
        lbl = tk.Label(top, image=tk_img)
        lbl.image = tk_img
        lbl.pack()

    def select_all(self):
        for rec in self.images:
            rec["var"].set(True)

    def invert_selection(self):
        for rec in self.images:
            rec["var"].set(not rec["var"].get())

    def save_selected(self):
        if not self.images:
            messagebox.showwarning("提示", "没有要保存的图片，请先获取")
            return
        os.makedirs(self.save_dir, exist_ok=True)
        total = 0
        for rec in self.images:
            if not rec["var"].get():
                continue
            ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            md = rec["md5"][:8]
            if self.save_orig_var.get():
                fname = os.path.join(self.save_dir, f"{ts}_{md}_orig.jpg")
                with open(fname, "wb") as f:
                    f.write(rec["raw"])
            if self.save_comp_var.get():
                comp_bytes, q = compress_image(rec["img"])
                fname = os.path.join(self.save_dir, f"{ts}_{md}_compressed.jpg")
                with open(fname, "wb") as f:
                    f.write(comp_bytes)
            total += 1
        messagebox.showinfo("完成", f"已保存 {total} 张（按选中项）")


if __name__ == "__main__":
    root = tk.Tk()
    root.geometry("1100x700")
    app = App(root)
    root.mainloop()
